import os
import torch
import torch.utils.data
import torchvision
import numpy as np

from data.apple_dataset import AppleDataset
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

import utility.utils as utils
import utility.transforms as T


######################################################
# Predict with either a Faster-RCNN or Mask-RCNN predictor
# using the MinneApple dataset
######################################################
def get_transform(train):
    transforms = []
    transforms.append(T.ToTensor())
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    return T.Compose(transforms)


def get_maskrcnn_model_instance(num_classes):
    # load an instance segmentation model pre-trained pre-trained on COCO
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False)

    # get number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # now get the number of input features for the mask classifier
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    # and replace the mask predictor with a new one
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)
    return model


def get_frcnn_model_instance(num_classes):
    # load an instance segmentation model pre-trained pre-trained on COCO
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)

    # get number of input features for the classifier
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model


def main(args):
    num_classes = 2
    device = args.device

    # Load the model from
    print("Loading model")
    # Create the correct model type
    if args.mrcnn:
        model = get_maskrcnn_model_instance(num_classes)
    else:
        model = get_frcnn_model_instance(num_classes)

    # # Load model parameters and keep on CPU
    # checkpoint = torch.load(args.weight_file, map_location=device)
    # # model.load_state_dict(checkpoint['model'], strict=False)
    # model.load_state_dict(checkpoint['model'] if 'model' in checkpoint else checkpoint)
    # model.eval()
    model = torch.load(args.weight_file, map_location=device)
    model.to(device)
    model.eval()
    print("Creating data loaders")
    dataset_test = AppleDataset(args.data_path, get_transform(train=False))
    data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1,
                                                   shuffle=False, num_workers=1,
                                                   collate_fn=utils.collate_fn)

    # Create output directory
    base_path = os.path.dirname(args.output_file )
    if not os.path.exists(base_path):
        os.makedirs(base_path)

    # Predict on bboxes on each image
    f = open(args.output_file, 'a')
    for image, targets in data_loader_test:
        image = list(img.to(device) for img in image)
        outputs = model(image)
        all_stacks = []
        for ii, output in enumerate(outputs):
            img_id = targets[ii]['image_id']
            img_name = data_loader_test.dataset.get_img_name(img_id)
            print(f"Predicting on image: {img_name}")

            boxes = output['boxes'].detach().cpu().numpy()
            scores = output['scores'].detach().cpu().numpy()

            # 生成该图片所有预测结果的组合数组
            im_names = np.repeat(img_name, len(boxes), axis=0)
            stacked = np.hstack((im_names.reshape(-1, 1), boxes.astype(int), scores.reshape(-1, 1)))
            all_stacks.append(stacked)

        # 整张图片的结果一次性写入
        if all_stacks:
            combined = np.vstack(all_stacks)
            np.savetxt(f, combined, fmt='%s', delimiter=',', newline='\n')


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='PyTorch Detection')
    parser.add_argument('--data_path',default='D:/Jupyter/pytorch/d2l-zh/pytorch/appppppp/detection2/test',  required=True, help='path to the data to predict on' )
    parser.add_argument('--output_file',default='./MinneApple-master/MinneApple-master/output', required=True, help='path where to write the prediction outputs')
    parser.add_argument('--weight_file', required=True ,default='./MinneApple-master/myapple.pth',help='path to the weight file' )
    parser.add_argument('--device', default='cuda', help='device to use. Either cpu or cuda')
    model = parser.add_mutually_exclusive_group(required=True)
    model.add_argument('--frcnn', action='store_true', help='use a Faster-RCNN model')
    model.add_argument('--mrcnn', action='store_true', help='use a Mask-RCNN model' )

    args = parser.parse_args()
    main(args)
# D:\anaconda3\envs\torch2\python.exe E:\download\MinneApple-master\MinneApple-master\MinneApple-master\detection_eval.py --data_path "D:/Jupyter/pytorch/d2l-zh/pytorch/appppppp/detection2/test" --output_file "./output" --weight_file "./myapple.pth" --device "cuda" --mrcnn
#
#
#


'''
这是一个基于PyTorch实现的苹果检测项目，包含预测和评估两个模块：
预测模块 (predict_rcnn.py)
支持Faster R-CNN和Mask R-CNN两种模型
主要功能： ✅ 加载预训练模型 ✅ 处理测试集图像 ✅ 输出预测框坐标和置信度 ✅ 结果保存为CSV格式（图片名, x1,y1,x2,y2, 分数）
评估模块 (detection_eval.py)
使用COCO评估指标
主要功能： ✅ 读取预测结果文件 ✅ 计算AP（平均精度）及其变体 ✅ 输出评估结果到scores.txt ✅ 支持不同尺度目标的评估（小/中/大目标）
两个模块通过命令行参数交互，典型使用流程：
用predict_rcnn.py生成检测结果
用detection_eval.py计算检测精度
最终输出包括检测框坐标和精度指标文件
 '''